# 在file_manager.py中添加以下功能
import pandas as pd
from pathlib import Path
from typing import Set
from config import settings
from config.exceptions import DataProcessError, BookProcessingError


class FileManager:
    @staticmethod
    def load_known_words_from_dir(data_dir: Path) -> Set[str]:
        """从目录中所有xlsx文件加载已知词汇"""
        known_words = set()

        try:
            # 遍历data目录下所有xlsx文件
            for xlsx_file in data_dir.glob('*.xlsx'):
                with pd.ExcelFile(xlsx_file) as excel:
                    # 读取所有sheet
                    for sheet_name in excel.sheet_names:
                        try:
                            df = pd.read_excel(
                                excel,
                                sheet_name=sheet_name,
                                usecols='A',
                                header=None,
                                names=['word'],
                                dtype={'word': str}
                            )
                            # 清理数据并转换为小写
                            words = df['word'].str.strip().dropna().str.lower()
                            known_words.update(words.astype(str))
                        except Exception as e:
                            raise DataProcessError(
                                f"读取文件 {xlsx_file.name} 工作表 {sheet_name} 失败: {str(e)}"
                            )
            return known_words
        except Exception as e:
            raise DataProcessError(f"加载已知词表失败: {str(e)}")

    # 修改原load_known_words函数
    @staticmethod
    def load_known_words() -> Set[str]:
        """加载所有已知词汇（包含目录和指定文件）"""
        try:
            # 加载目录词汇
            data_dir = settings.BASE_DIR / 'data'
            dir_words = FileManager.load_known_words_from_dir(data_dir)

            # 加载指定文件词汇
            file_words = set()
            if settings.KNOWN_WORDS_PATH.exists():
                df = pd.read_excel(
                    settings.KNOWN_WORDS_PATH,
                    sheet_name=None,  # 读取所有sheet
                    usecols='A',
                    header=None,
                    names=['word'],
                    dtype={'word': str}
                )
                for sheet in df.values():
                    words = sheet['word'].str.strip().dropna().str.lower()
                    file_words.update(words.astype(str))

            # 合并结果
            return dir_words.union(file_words)
        except Exception as e:
            raise DataProcessError(f"综合加载失败: {str(e)}")



    @staticmethod
    def save_vocabulary(data, output_path):
        try:
            writer = pd.ExcelWriter(output_path, engine='openpyxl')
            for sheet_name, df in data.items():
                df.to_excel(writer, sheet_name=sheet_name, index=False)
            writer.close()
        except Exception as e:
            raise BookProcessingError(f"文件保存失败: {str(e)}")